import gym
import pygame
from gym.utils.play import play

'''键盘控制智能体：'''
play(gym.make('Pong-v0'))

# 定义键到动作的映射
mapping = {
    (pygame.K_LEFT,): 0,  # 按下左箭头键 -> 执行动作 0
    (pygame.K_RIGHT,): 1  # 按下右箭头键 -> 执行动作 1
}

# 启动环境并使用自定义按键映射
play(gym.make("CartPole-v0"), keys_to_action=mapping)

# 如果你想在游戏过程中绘制一些实时统计数据（如奖励曲线），可以使用 gym.utils.play.PlayPlot 工具。例如，以下代码绘制过去 5 秒钟的奖励变化：
from gym.utils.play import PlayPlot


# 定义回调函数，返回奖励
def callback(obs_t, obs_tp1, action, rew, done, info):
    return [rew]  # 返回奖励值


# 初始化绘图工具
plotter = PlayPlot(callback, 30 * 5, ["reward"])  # 绘制最近 5 秒的奖励
# 启动环境并应用绘图工具
env = gym.make("Pong-v0")
play(env, callback=plotter.callback)
